import sqlite3
import os
from tabulate import tabulate
import pandas as pd

def query_database():
    # 获取数据库路径
    base_dir = os.path.dirname(os.path.abspath(__file__))
    db_path = os.path.join(base_dir, 'data', 'predictions.db')
    
    if not os.path.exists(db_path):
        print(f"数据库文件不存在: {db_path}")
        return
    
    # 连接到数据库
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        print(f"成功连接到数据库: {db_path}")
        
        # 获取所有表名
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
        tables = [table[0] for table in cursor.fetchall()]
        print(f"\n数据库中的表: {tables}")
        
        # 查询模型性能表
        try:
            print("\n== 模型性能表 ==")
            cursor.execute("SELECT * FROM model_performance")
            rows = cursor.fetchall()
            
            if not rows:
                print("模型性能表中没有数据")
            else:
                # 获取列名
                cursor.execute("PRAGMA table_info(model_performance)")
                columns = [col[1] for col in cursor.fetchall()]
                
                # 使用pandas美化输出
                df = pd.DataFrame(rows, columns=columns)
                print(tabulate(df, headers='keys', tablefmt='psql', showindex=False))
                print(f"总记录数: {len(rows)}")
        except sqlite3.Error as e:
            print(f"查询模型性能表时出错: {e}")
        
        # 查询预测记录表
        try:
            print("\n== 预测记录表 ==")
            cursor.execute("SELECT id, user_id, prediction, probability, timestamp FROM predictions")
            rows = cursor.fetchall()
            
            if not rows:
                print("预测记录表中没有数据")
            else:
                # 获取列名 (只显示部分列)
                columns = ['id', 'user_id', 'prediction', 'probability', 'timestamp']
                
                # 使用pandas美化输出
                df = pd.DataFrame(rows, columns=columns)
                print(tabulate(df, headers='keys', tablefmt='psql', showindex=False))
                print(f"总记录数: {len(rows)}")
                
                # 按用户ID统计
                print("\n按用户ID统计的预测次数:")
                cursor.execute("SELECT user_id, COUNT(*) FROM predictions GROUP BY user_id")
                user_stats = cursor.fetchall()
                if user_stats:
                    user_df = pd.DataFrame(user_stats, columns=['用户ID', '预测次数'])
                    print(tabulate(user_df, headers='keys', tablefmt='psql', showindex=False))
        except sqlite3.Error as e:
            print(f"查询预测记录表时出错: {e}")
            
    except sqlite3.Error as e:
        print(f"连接数据库时出错: {e}")
    finally:
        if conn:
            conn.close()
            print("\n数据库连接已关闭")

if __name__ == "__main__":
    try:
        query_database()
    except Exception as e:
        print(f"发生错误: {e}") 